from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, SkipExhausted


class CounterIterator(DatasetIterator[int]):
    def __init__(self, start: int, stop: int | None = None, step: int = 1):
        self._start = start
        self._stop = stop
        self._step = step

    def __next__(self) -> int:
        if self._stop is not None and self._start >= self._stop:
            raise StopIteration
        index = self._start
        self._start += self._step
        return index

    def __skip__(self, samples_n: int):
        if self._stop is not None:
            remain_samples_n = (self._stop - self._start) // self._step + 1 \
                if self._stop > self._start else 0
            if samples_n > remain_samples_n:
                self._start += remain_samples_n * self._step
                raise SkipExhausted(remain_samples_n)
        self._start += samples_n * self._step

    def __pause__(self) -> int:
        return self._start


class CounterDataset(Dataset[int]):
    def __init__(self, start: int = 0, stop: int | None = None, step: int = 1):
        self._start = start
        self._stop = stop
        self._step = step

    def __len__(self) -> int:
        if self._stop is None:
            raise TypeError
        return (self._stop - self._start) // self._step + 1 \
            if self._stop > self._start else 0

    def __getitem__(self, index: int, /) -> int:
        if index < 0:
            raise IndexError
        if index >= len(self):
            raise IndexError
        return self._start + index * self._step

    def __iter__(self) -> DatasetIterator[int]:
        return CounterIterator(self._start)

    def __resume__(self, state: int | None = None, /) -> DatasetIterator[int]:
        return CounterIterator(self._start if state is None else state)
